{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "s4eK-unQlKTF"
   },
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pinecone-io/examples/blob/master/learn/generation/llama-index/llama-index-intro.ipynb) [![Open nbviewer](https://raw.githubusercontent.com/pinecone-io/examples/master/assets/nbviewer-shield.svg)](https://nbviewer.org/github/pinecone-io/examples/blob/master/learn/generation/llama-index/llama-index-intro.ipynb)\n",
    "\n",
    "# Llama-Index with Pinecone\n",
    "\n",
    "**Note:** This notebook is built to run end-to-end in Google Colab. Users might run into issues when running locally, depending on their local environment setup.\n",
    "\n",
    "In this notebook, we will demo how to use the `llama-index` (previously GPT-index) library with Pinecone for semantic search. This notebook is an introduction and does not cover the more advanced features of `llama-index`. You will find these in future releases in the [Pinecone examples repo](https://github.com/pinecone-io/examples/tree/master/learn/generation).\n",
    "\n",
    "We will start by installing the necessary libraries and initializing Pinecone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "YenjP38jtRdE"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.2.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.3.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install -qU llama-index==0.9.29 datasets==2.16.1 pinecone-client==3.1.0 openai==1.7.1 transformers==4.36.2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "KjeiHqOIoGnI"
   },
   "source": [
    "We can go ahead and load the **SQuAD** dataset, which contains questions and answer pairs from Wikipedia articles. We'll then convert the dataset into a pandas DataFrame and keep only the unique 'context' fields, which are the text passages that the questions are based on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 244
    },
    "id": "TlKkc8Jbqiin",
    "outputId": "34a6aa74-477e-45bf-fede-4d89603c291b"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>context</th>\n",
       "      <th>title</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5733be284776f41900661182</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>5733bf84d058e614000b61be</td>\n",
       "      <td>As at most other universities, Notre Dame's st...</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>5733bed24776f41900661188</td>\n",
       "      <td>The university is the major seat of the Congre...</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>5733a6424776f41900660f51</td>\n",
       "      <td>The College of Engineering was established in ...</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>5733a70c4776f41900660f64</td>\n",
       "      <td>All of Notre Dame's undergraduate students are...</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                          id  \\\n",
       "0   5733be284776f41900661182   \n",
       "5   5733bf84d058e614000b61be   \n",
       "10  5733bed24776f41900661188   \n",
       "15  5733a6424776f41900660f51   \n",
       "20  5733a70c4776f41900660f64   \n",
       "\n",
       "                                              context  \\\n",
       "0   Architecturally, the school has a Catholic cha...   \n",
       "5   As at most other universities, Notre Dame's st...   \n",
       "10  The university is the major seat of the Congre...   \n",
       "15  The College of Engineering was established in ...   \n",
       "20  All of Notre Dame's undergraduate students are...   \n",
       "\n",
       "                       title  \n",
       "0   University_of_Notre_Dame  \n",
       "5   University_of_Notre_Dame  \n",
       "10  University_of_Notre_Dame  \n",
       "15  University_of_Notre_Dame  \n",
       "20  University_of_Notre_Dame  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "data = load_dataset('squad', split='train')\n",
    "data = data.to_pandas()[['id', 'context', 'title']]\n",
    "data.drop_duplicates(subset='context', keep='first', inplace=True)\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZT6Z5gW4oTkG",
    "outputId": "c4797ae0-4091-4b8e-fa63-293174deb073"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "18891"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "IK90FFhf42hd"
   },
   "source": [
    "This code transforms our DataFrame into a list of Document objects, ready for indexing with llama_index. Each document contains the text passage, a unique id, and an extra field for the article title."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uXOYMGfLtO-Z",
    "outputId": "62466814-5797-40f4-c606-a163c44581e4"
   },
   "outputs": [],
   "source": [
    "from llama_index import Document\n",
    "\n",
    "docs = []\n",
    "\n",
    "for i, row in data.iterrows():\n",
    "    docs.append(Document(\n",
    "        text=row['context'],\n",
    "        doc_id=row['id'],\n",
    "        extra_info={'title': row['title']}\n",
    "    ))\n",
    "docs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hG81jmfCobYF",
    "outputId": "e8813daa-98ad-4f37-a0dd-560046a1073a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "18891"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(docs)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "uQ4d9UlS5OeU"
   },
   "source": [
    "Here, we're setting up the OpenAI API key and initializing a `SimpleNodeParser`. This parser processes our list of `Document` objects into 'nodes', which are the basic units that `llama_index` uses for indexing and querying. The first node is displayed below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "1MSXio_atoSq"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['OPENAI_API_KEY'] = '<your OpenAI API key>'  # platform.openai.com"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gKTwFpkmte7o",
    "outputId": "f3774aa6-4fca-44a8-ae1f-58063eae50c6"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Node(text='Architecturally, the school has a Catholic character. Atop the Main Building\\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend \"Venite Ad Me Omnes\". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.', doc_id='8968aa3c-abbe-4fe0-ad91-be0298830542', embedding=None, doc_hash='ba8eb7d843763d4f7e7f71f0be64c70c4597fd712bfcce77007fc4c821b8d1b6', extra_info={'title': 'University_of_Notre_Dame'}, node_info={'start': 0, 'end': 695}, relationships={<DocumentRelationship.SOURCE: '1'>: '5733be284776f41900661182'})"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from llama_index.node_parser import SimpleNodeParser\n",
    "\n",
    "parser = SimpleNodeParser()\n",
    "\n",
    "nodes = parser.get_nodes_from_documents(docs)\n",
    "nodes[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "LQL4X2_WodU2",
    "outputId": "34d05a8e-1c39-4248-e4d6-25433fd31b58"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "18892"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(nodes)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "JQrAcAdZ1fS9"
   },
   "source": [
    "### Indexing in Pinecone\n",
    "\n",
    "Pinecone is a managed vector database service designed for machine learning applications. We're using it in this context to store and retrieve embeddings generated by our language model, enabling efficient and scalable semantic similarity-based search.\n",
    "\n",
    "We initialize Pinecone with the relevant API key and environment that we [get for **free** in the console](https://app.pinecone.io/), then create a new index. The index has a dimension of 1536 and uses cosine similarity, which is the recommended metric for comparing vectors produced by the `text-embedding-ada-002` model we'll be using."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pinecone import Pinecone\n",
    "from pinecone import ServerlessSpec\n",
    "\n",
    "\n",
    "# find API key in console at app.pinecone.io\n",
    "os.environ['PINECONE_API_KEY'] = '<your API key here>'\n",
    "\n",
    "# initialize connection to pinecone\n",
    "pc = Pinecone(api_key=os.getenv('PINECONE_API_KEY'))\n",
    "\n",
    "\n",
    "# create the index if it does not exist already\n",
    "index_name = 'llama-index-intro'\n",
    "existing_indexes = [i.get('name') for i in pc.list_indexes()]\n",
    "\n",
    "if index_name not in existing_indexes:\n",
    "    pc.create_index(\n",
    "        name=index_name,\n",
    "        dimension=1536,\n",
    "        metric='cosine',\n",
    "        spec=ServerlessSpec(cloud='aws', region='us-west-2')\n",
    "    )\n",
    "\n",
    "# connect to the index\n",
    "pinecone_index = pc.Index(index_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "r4KQXwLh645t"
   },
   "source": [
    "Here, we're initializing a `PineconeVectorStore` with our previously created Pinecone index. This object will serve as the storage and retrieval interface for our document embeddings in Pinecone's vector database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "AO4eAvSn4t-R"
   },
   "outputs": [],
   "source": [
    "from llama_index.vector_stores import PineconeVectorStore\n",
    "\n",
    "# we can select a namespace (acts as a partition in an index)\n",
    "namespace = '' # default namespace\n",
    "\n",
    "vector_store = PineconeVectorStore(pinecone_index=pinecone_index)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "qotocu00ZjfF"
   },
   "source": [
    "Next we initialize the `GPTVectorStoreIndex` with our list of `Document` objects, using the `PineconeVectorStore` as storage and `OpenAIEmbedding` model for embeddings.\n",
    "\n",
    "`StorageContext` is used to configure the storage setup, and `ServiceContext` sets up the embedding model. The `GPTVectorStoreIndex` handles the indexing and querying process, making use of the provided storage and service contexts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "V64_Y5FZqJ8w"
   },
   "outputs": [],
   "source": [
    "from llama_index import GPTVectorStoreIndex, StorageContext, ServiceContext\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "# setup our storage (vector db)\n",
    "storage_context = StorageContext.from_defaults(\n",
    "    vector_store=vector_store\n",
    ")\n",
    "# setup the index/query process, ie the embedding model (and completion if used)\n",
    "embed_model = OpenAIEmbedding(model='text-embedding-ada-002', embed_batch_size=100)\n",
    "service_context = ServiceContext.from_defaults(embed_model=embed_model)\n",
    "\n",
    "index = GPTVectorStoreIndex.from_documents(\n",
    "    docs, storage_context=storage_context,\n",
    "    service_context=service_context\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "Z7veOHn3Zyv1"
   },
   "source": [
    "Finally we can build a query engine from the `index` we build and use this engine to perform a query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4loa7g_y5x2v",
    "outputId": "e87ab48f-f641-46a1-f06b-be2d5daf4a0b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The College of Engineering was established in 1920.\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine()\n",
    "res = query_engine.query(\"in what year was the college of engineering established at the University of Notre Dame?\")\n",
    "print(res)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "xQEzZJqWaAgS"
   },
   "source": [
    "That's our quick intro to using Llama-index and Pinecone! Once we're done testing the system we should delete the Pinecone index to save resources:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "z89CVJMGaJ4j"
   },
   "outputs": [],
   "source": [
    "pc.delete_index(index_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "DWtX3Nxl5V71"
   },
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
